#include <chrono>
#include <fstream>
#include <iostream>
#include <openssl/aes.h>
#include <openssl/evp.h>
#include <string>

#define EVP_MAX_TAG_LENGTH 16

size_t get_file_size(const std::string& filename)
{
    std::ifstream ifs(filename.c_str(), std::ios::in | std::ios::binary);
    if (!ifs.good()) {
        std::cout << "Failed to open the file \"" << filename << "\"" << std::endl;
        return -1;
    }
    ifs.seekg(0, std::ios::end);
    size_t size = (size_t)ifs.tellg();
    return size;
}

// 加密函数
bool encryptFile(const std::string& inputFile, const std::string& outputFile, const unsigned char* key, const unsigned char* iv)
{
    const int BUFFER_SIZE = get_file_size(inputFile);

    // 打开输入文件
    std::ifstream fin(inputFile, std::ios::binary);
    if (!fin) {
        std::cerr << "Failed to open input file: " << inputFile << std::endl;
        return false;
    }

    // 打开输出文件
    std::ofstream fout(outputFile, std::ios::binary);
    if (!fout) {
        std::cerr << "Failed to open output file: " << outputFile << std::endl;
        fin.close();
        return false;
    }

    // 创建 AES 上下文
    EVP_CIPHER_CTX* ctx = EVP_CIPHER_CTX_new();
    if (!ctx) {
        std::cerr << "Failed to create AES context" << std::endl;
        fin.close();
        fout.close();
        return false;
    }

    // 初始化 AES 加密操作
    if (EVP_EncryptInit_ex(ctx, EVP_aes_128_gcm(), NULL, NULL, NULL) != 1) {
        std::cerr << "Failed to initialize AES encryption" << std::endl;
        EVP_CIPHER_CTX_free(ctx);
        fin.close();
        fout.close();
        return false;
    }

    // 设置 AES 密钥和 IV
    if (EVP_EncryptInit_ex(ctx, NULL, NULL, key, iv) != 1) {
        std::cerr << "Failed to set AES key and IV" << std::endl;
        EVP_CIPHER_CTX_free(ctx);
        fin.close();
        fout.close();
        return false;
    }

    // 读取和加密输入文件内容
    char* buffer = new char[BUFFER_SIZE];
    int bytesRead;
    int bytesWritten;
    std::chrono::_V2::steady_clock::time_point pos1;
    while ((bytesRead = fin.read(buffer, BUFFER_SIZE).gcount()) > 0) {
        pos1 = std::chrono::steady_clock::now();
        if (EVP_EncryptUpdate(ctx, reinterpret_cast<unsigned char*>(buffer), &bytesWritten,
                reinterpret_cast<const unsigned char*>(buffer), bytesRead)
            != 1) {
            std::cerr << "AES encryption error" << std::endl;
            EVP_CIPHER_CTX_free(ctx);
            delete[] buffer;
            fin.close();
            fout.close();
            return false;
        }
    }

    // 完成 AES 加密操作
    unsigned char* finalBuffer = new unsigned char[BUFFER_SIZE];
    int finalBytesWritten;
    if (EVP_EncryptFinal_ex(ctx, finalBuffer, &finalBytesWritten) != 1) {
        std::cerr << "Failed to finalize AES encryption" << std::endl;
        EVP_CIPHER_CTX_free(ctx);
        delete[] buffer;
        delete[] finalBuffer;
        fin.close();
        fout.close();
        return false;
    }

    // 获取认证标签
    unsigned char tag[EVP_MAX_TAG_LENGTH];
    int tagLength;
    if (EVP_CIPHER_CTX_ctrl(ctx, EVP_CTRL_GCM_GET_TAG, EVP_MAX_TAG_LENGTH, tag) != 1) {
        std::cerr << "Failed to get authentication tag" << std::endl;
        EVP_CIPHER_CTX_free(ctx);
        delete[] buffer;
        delete[] finalBuffer;
        fin.close();
        fout.close();
        return false;
    }
    std::chrono::_V2::steady_clock::time_point pos2 = std::chrono::steady_clock::now();

    fout.write(buffer, bytesWritten);
    fout.write(reinterpret_cast<const char*>(finalBuffer), finalBytesWritten);
    fout.write(reinterpret_cast<const char*>(tag), EVP_MAX_TAG_LENGTH);

    std::chrono::duration<double> dur1 = pos2 - pos1;
    std::cout << "Decrypt ["
              << "] bytes data cost time: " << dur1.count() << "s" << std::endl;

    // std::cout << "Succeed to decrypt data from \"" << ENCRYPT_FILE_NAME << "\", data size: " << full_size - SGX_AESGCM_MAC_SIZE << std::endl;

    // 释放资源
    EVP_CIPHER_CTX_free(ctx);
    delete[] buffer;
    delete[] finalBuffer;
    fin.close();
    fout.close();

    std::cout << "Encryption completed successfully" << std::endl;
    return true;
}

int main()
{
    std::string inputFile = "/root/sgx/linux-sgx/src/SealUnseal/data_need_encrypt";
    std::string outputFile = "/root/sgx/linux-sgx/src/SealUnseal/data_encrypted_outsgx";
    const unsigned char key[16] = { 0x11, 0x22, 0x33, 0x44, 0x55, 0x66, 0x77, 0x88, 0x99, 0x00, 0x11, 0x22, 0x33, 0x44, 0x55, 0x66 }; // 16 字节的密钥
    const unsigned char iv[12] = { 0x00, 0x11, 0x22, 0x11, 0x22, 0x33, 0x44, 0x55, 0x66, 0x77, 0x88, 0x99 }; // 12 字节的 IV

    if (encryptFile(inputFile, outputFile, key, iv)) {
        std::cout << "Encryption successful" << std::endl;
    } else {
        std::cerr << "Encryption failed" << std::endl;
    }

    return 0;
}
